# multi_and_mitm.py —— 双重/三重加密与中间相遇攻击（MITM）
import saes

# --- 双重加密（32bit 密钥：K1||K2，各16bit）---
def double_encrypt_block16(p16: int, key32: int) -> int:
    k1 = (key32 >> 16) & 0xFFFF
    k2 = key32 & 0xFFFF
    return saes.encrypt_block16(saes.encrypt_block16(p16, k1), k2)

def double_decrypt_block16(c16: int, key32: int) -> int:
    k1 = (key32 >> 16) & 0xFFFF
    k2 = key32 & 0xFFFF
    return saes.decrypt_block16(saes.decrypt_block16(c16, k2), k1)

# --- 三重加密：两种模式任选其一 ---
# (A) 2-key EDE： E_{K1}( D_{K2}( E_{K1}(P) ) )
def triple_2key_EDE_encrypt(p16: int, k1: int, k2: int) -> int:
    return saes.encrypt_block16(saes.decrypt_block16(saes.encrypt_block16(p16, k1), k2), k1)

def triple_2key_EDE_decrypt(c16: int, k1: int, k2: int) -> int:
    return saes.decrypt_block16(saes.encrypt_block16(saes.decrypt_block16(c16, k1), k2), k1)

# (B) 3-key EEE： E_{K3}( E_{K2}( E_{K1}(P) ) )
def triple_3key_EEE_encrypt(p16: int, k1: int, k2: int, k3: int) -> int:
    return saes.encrypt_block16(saes.encrypt_block16(saes.encrypt_block16(p16, k1), k2), k3)

def triple_3key_EEE_decrypt(c16: int, k1: int, k2: int, k3: int) -> int:
    return saes.decrypt_block16(saes.decrypt_block16(saes.decrypt_block16(c16, k3), k2), k1)

# --- 中间相遇攻击（给定一对 P, C，暴力 2^16 键空间；可扩展到多对筛选）---
def mitm_attack(P: int, C: int):
    # 构造 E_{k1}(P) 哈希表
    forward = {}
    for k1 in range(0x10000):
        mid = saes.encrypt_block16(P, k1)
        forward[mid] = forward.get(mid, []) + [k1]

    # 遍历 k2 比较 D_{k2}(C)
    candidates = []
    for k2 in range(0x10000):
        mid = saes.decrypt_block16(C, k2)
        if mid in forward:
            for k1 in forward[mid]:
                candidates.append((k1, k2))
    return candidates  # 可能有多个；可用第二组 (P2,C2) 继续筛

if __name__ == "__main__":
    P = 0x1234
    K1, K2 = 0xAAAA, 0x5555
    C = double_encrypt_block16(P, (K1<<16)|K2)
    print(f"Double: P={P:04X}, C={C:04X}")
    pairs = mitm_attack(P, C)
    print("Found pairs (k1,k2):", [(f"{a:04X}", f"{b:04X}") for a,b in pairs[:5]])
